/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.benchmark.perfetto

import android.os.Build
import androidx.annotation.RequiresApi
import androidx.annotation.RestrictTo
import androidx.benchmark.Arguments
import androidx.benchmark.Shell
import androidx.benchmark.VirtualFile
import perfetto.protos.AndroidPowerConfig
import perfetto.protos.DataSourceConfig
import perfetto.protos.FtraceConfig
import perfetto.protos.HeapprofdConfig
import perfetto.protos.MeminfoCounters
import perfetto.protos.PerfEventConfig
import perfetto.protos.PerfEvents
import perfetto.protos.ProcessStatsConfig
import perfetto.protos.SysStatsConfig
import perfetto.protos.TraceConfig
import perfetto.protos.TraceConfig.BufferConfig
import perfetto.protos.TraceConfig.BufferConfig.FillPolicy
import perfetto.protos.TrackEventConfig

/**
 * Configuration for Perfetto trace recording.
 *
 * For more info, see https://perfetto.dev/docs/concepts/config
 */
@ExperimentalPerfettoCaptureApi
sealed class PerfettoConfig(internal val isTextProto: Boolean) {
    @RequiresApi(23) internal abstract fun writeTo(virtualFile: VirtualFile)

    /**
     * Binary representation of a Perfetto config proto.
     *
     * This can be generated by a proto library, together with the definition here:
     */
    class Binary(val bytes: ByteArray) : PerfettoConfig(isTextProto = false) {
        @RequiresApi(23)
        override fun writeTo(virtualFile: VirtualFile) {
            virtualFile.writeBytes(bytes)
        }
    }

    /**
     * TextProto representation of a Perfetto config.
     *
     * This can be generated with https://ui.perfetto.dev/#!/record/
     *
     * Note: this format is not recommended for long term use - the [Binary] proto representation is
     * more likely to remain stable over time, across Perfetto/Android OS versions. For more
     * information, see
     * [the Perfetto documentation](https://perfetto.dev/docs/concepts/config#pbtx-vs-binary-format).
     */
    class Text(val text: String) : PerfettoConfig(isTextProto = true) {
        @RequiresApi(23)
        override fun writeTo(virtualFile: VirtualFile) {
            virtualFile.writeText(text)
        }
    }

    /** Benchmark defined config for perfetto trace capture, used by benchmark/macrobenchmark. */
    @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
    class Benchmark(
        private val appTagPackages: List<String>,
        private val useStackSamplingConfig: Boolean,
    ) : PerfettoConfig(isTextProto = false) {
        @RequiresApi(23)
        override fun writeTo(virtualFile: VirtualFile) {
            val stackSamplingConfig =
                if (useStackSamplingConfig) {
                    Arguments.profiler?.config(appTagPackages)
                } else {
                    null
                }
            virtualFile.writeBytes(
                perfettoConfig(
                        atraceApps =
                            if (Build.VERSION.SDK_INT <= 28 || appTagPackages.isEmpty()) {
                                appTagPackages
                            } else {
                                listOf("*")
                            },
                        stackSamplingConfig = stackSamplingConfig,
                    )
                    .validateAndEncode()
            )
        }
    }

    /**
     * Only used by group-internal tests
     *
     * Most minimal config possible
     */
    @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
    class MinimalTest(private val appTagPackages: List<String>) :
        PerfettoConfig(isTextProto = false) {
        @RequiresApi(23)
        override fun writeTo(virtualFile: VirtualFile) {
            virtualFile.writeBytes(
                configOf(listOf(minimalAtraceDataSource(atraceApps = appTagPackages)))
                    .validateAndEncode()
            )
        }
    }
}

private fun minimalAtraceDataSource(atraceApps: List<String>) =
    TraceConfig.DataSource(
        config =
            DataSourceConfig(
                name = "linux.ftrace",
                target_buffer = 0,
                ftrace_config =
                    FtraceConfig(
                        ftrace_events = emptyList(),
                        atrace_categories = emptyList(),
                        atrace_apps = atraceApps,
                        compact_sched = null,
                    ),
            )
    )

private fun ftraceDataSource(atraceApps: List<String>) =
    TraceConfig.DataSource(
        config =
            DataSourceConfig(
                name = "linux.ftrace",
                target_buffer = 0,
                ftrace_config =
                    FtraceConfig(
                        ftrace_events =
                            listOf(
                                // We need to do process tracking to ensure kernel ftrace events
                                // targeted at short-lived
                                // threads are associated correctly
                                "task/task_newtask",
                                "task/task_rename",
                                "sched/sched_process_exit",
                                "sched/sched_process_free",

                                // Memory events
                                "mm_event/mm_event_record",
                                "kmem/rss_stat",
                                "kmem/ion_heap_shrink",
                                "kmem/ion_heap_grow",
                                "ion/ion_stat",
                                "oom/oom_score_adj_update",

                                // Disk I/O
                                "disk",
                                "ufs/ufshcd_clk_gating",

                                // Old (kernel) LMK
                                "lowmemorykiller/lowmemory_kill",
                            ),
                        atrace_categories =
                            listOf(
                                    AtraceTag.ActivityManager,
                                    AtraceTag.Aidl,
                                    AtraceTag.Audio,
                                    AtraceTag.BinderDriver,
                                    AtraceTag.Camera,
                                    AtraceTag.Dalvik,
                                    AtraceTag.Frequency,
                                    AtraceTag.Graphics,
                                    AtraceTag.HardwareModules,
                                    AtraceTag.Idle,
                                    AtraceTag.Input,
                                    AtraceTag.MemReclaim,
                                    AtraceTag.Power,
                                    AtraceTag.Resources,
                                    AtraceTag.Scheduling,
                                    AtraceTag.Synchronization,
                                    AtraceTag.View,
                                    AtraceTag.WindowManager,
                                    // "webview" not included to workaround b/190743595
                                    // "memory" not included as some Q devices requiring
                                    // ftrace_event
                                    // configuration directly to collect this data. See b/171085599
                                )
                                .filter {
                                    // filter to only supported tags on unrooted build
                                    // TODO: use root-only tags as needed
                                    it.supported(api = Build.VERSION.SDK_INT, rooted = false)
                                }
                                .map { it.tag },
                        atrace_apps = atraceApps,
                        compact_sched = FtraceConfig.CompactSchedConfig(enabled = true),
                    ),
            )
    )

private fun processStatsDataSource(
    stackSamplingConfig: StackSamplingConfig?
): TraceConfig.DataSource {
    return TraceConfig.DataSource(
        config =
            DataSourceConfig(
                name = "linux.process_stats",
                target_buffer = 1,
                process_stats_config =
                    ProcessStatsConfig(
                        proc_stats_poll_ms = stackSamplingConfig?.frequency?.toInt() ?: 10000,
                        // This flag appears to be unreliable on API 29 unbundled perfetto, so to
                        // avoid very
                        // frequent proc stats polling to name processes correctly, we currently use
                        // unbundled
                        // perfetto on API 29, even though the bundled version exists. (b/218668335)
                        scan_all_processes_on_start = true,
                    ),
            )
    )
}

private val PACKAGE_LIST_DATASOURCE =
    TraceConfig.DataSource(
        config = DataSourceConfig(name = "android.packages_list", target_buffer = 1)
    )

private val LINUX_SYS_STATS_DATASOURCE =
    TraceConfig.DataSource(
        config =
            DataSourceConfig(
                name = "linux.sys_stats",
                target_buffer = 1,
                sys_stats_config =
                    SysStatsConfig(
                        meminfo_period_ms = 1000,
                        meminfo_counters =
                            listOf(
                                MeminfoCounters.MEMINFO_MEM_TOTAL,
                                MeminfoCounters.MEMINFO_MEM_FREE,
                                MeminfoCounters.MEMINFO_MEM_AVAILABLE,
                                MeminfoCounters.MEMINFO_BUFFERS,
                                MeminfoCounters.MEMINFO_CACHED,
                                MeminfoCounters.MEMINFO_SWAP_CACHED,
                                MeminfoCounters.MEMINFO_ACTIVE,
                                MeminfoCounters.MEMINFO_INACTIVE,
                                MeminfoCounters.MEMINFO_ACTIVE_ANON,
                                MeminfoCounters.MEMINFO_INACTIVE_ANON,
                                MeminfoCounters.MEMINFO_ACTIVE_FILE,
                                MeminfoCounters.MEMINFO_INACTIVE_FILE,
                                MeminfoCounters.MEMINFO_UNEVICTABLE,
                                MeminfoCounters.MEMINFO_SWAP_TOTAL,
                                MeminfoCounters.MEMINFO_SWAP_FREE,
                                MeminfoCounters.MEMINFO_DIRTY,
                                MeminfoCounters.MEMINFO_WRITEBACK,
                                MeminfoCounters.MEMINFO_ANON_PAGES,
                                MeminfoCounters.MEMINFO_MAPPED,
                                MeminfoCounters.MEMINFO_SHMEM,
                            ),
                    ),
            )
    )

private val ANDROID_POWER_DATASOURCE =
    TraceConfig.DataSource(
        config =
            DataSourceConfig(
                name = "android.power",
                android_power_config =
                    AndroidPowerConfig(
                        battery_poll_ms = 250,
                        battery_counters =
                            listOf(
                                AndroidPowerConfig.BatteryCounters.BATTERY_COUNTER_CAPACITY_PERCENT,
                                AndroidPowerConfig.BatteryCounters.BATTERY_COUNTER_CHARGE,
                                AndroidPowerConfig.BatteryCounters.BATTERY_COUNTER_CURRENT,
                            ),
                        collect_power_rails = true,
                    ),
            )
    )

/** A Perfetto data source to enable stack sampling. */
private fun stackSamplingSource(config: StackSamplingConfig): List<TraceConfig.DataSource> {
    val sources = mutableListOf<TraceConfig.DataSource>()
    sources +=
        TraceConfig.DataSource(
            config =
                DataSourceConfig(
                    name = "linux.perf",
                    target_buffer = 1,
                    perf_event_config =
                        PerfEventConfig(
                            timebase =
                                PerfEvents.Timebase(
                                    counter = PerfEvents.Counter.SW_CPU_CLOCK,
                                    frequency = config.frequency,
                                    timestamp_clock = PerfEvents.PerfClock.PERF_CLOCK_MONOTONIC,
                                ),
                            callstack_sampling =
                                PerfEventConfig.CallstackSampling(
                                    scope =
                                        PerfEventConfig.Scope(target_cmdline = config.packageNames)
                                ),
                            kernel_frames = false,
                        ),
                )
        )
    sources +=
        TraceConfig.DataSource(
            config =
                DataSourceConfig(
                    name = "linux.perf",
                    target_buffer = 1,
                    perf_event_config =
                        PerfEventConfig(
                            timebase =
                                PerfEvents.Timebase(
                                    tracepoint = PerfEvents.Tracepoint(name = "sched_switch"),
                                    period = 1,
                                ),
                            callstack_sampling =
                                PerfEventConfig.CallstackSampling(
                                    scope =
                                        PerfEventConfig.Scope(target_cmdline = config.packageNames)
                                ),
                            kernel_frames = false,
                        ),
                )
        )
    if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
        // https://perfetto.dev/docs/reference/trace-config-proto#HeapprofdConfig
        sources +=
            TraceConfig.DataSource(
                config =
                    DataSourceConfig(
                        name = "android.heapprofd",
                        heapprofd_config =
                            HeapprofdConfig(
                                shmem_size_bytes = 8388608,
                                sampling_interval_bytes = 2048,
                                block_client = true,
                                process_cmdline = config.packageNames,
                                heaps =
                                    listOf(
                                        "com.android.art" // Java Heaps
                                    ),
                                continuous_dump_config =
                                    HeapprofdConfig.ContinuousDumpConfig(
                                        dump_phase_ms = 0,
                                        dump_interval_ms = 500, // ms
                                    ),
                            ),
                    )
            )
    }
    return sources
}

// reduce timeout to reduce trace capture overhead when devices have data source issues
// See b/323601788 and b/307649002.
internal const val PERFETTO_DATA_SOURCE_STOP_TIMEOUT_MS = 2500

private fun configOf(dataSources: List<TraceConfig.DataSource>) =
    TraceConfig(
        buffers =
            listOf(
                BufferConfig(size_kb = 32768, FillPolicy.RING_BUFFER),
                BufferConfig(size_kb = 4096, FillPolicy.RING_BUFFER),
            ),
        data_sources = dataSources,
        // periodically dump to file, so we don't overrun our ring buffer
        // buffers are expected to be big enough for 5 seconds, so conservatively set 2.5 dump
        write_into_file = true,
        file_write_period_ms = 2500,

        // multiple of file_write_period_ms, enables trace processor to work in batches
        flush_period_ms = 5000,
        data_source_stop_timeout_ms = PERFETTO_DATA_SOURCE_STOP_TIMEOUT_MS,
    )

/**
 * Config for perfetto.
 *
 * Eventually, this should be more configurable.
 */
@RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
internal fun perfettoConfig(
    atraceApps: List<String>,
    stackSamplingConfig: StackSamplingConfig?,
): TraceConfig {
    val dataSources =
        mutableListOf(
            ftraceDataSource(atraceApps),
            processStatsDataSource(stackSamplingConfig),
            PACKAGE_LIST_DATASOURCE,
            LINUX_SYS_STATS_DATASOURCE,
            ANDROID_POWER_DATASOURCE,
            TraceConfig.DataSource(DataSourceConfig("android.gpu.memory")),
            TraceConfig.DataSource(DataSourceConfig("android.surfaceflinger.frame")),
            TraceConfig.DataSource(DataSourceConfig("android.surfaceflinger.frametimeline")),
            TraceConfig.DataSource(
                DataSourceConfig(
                    "track_event",
                    track_event_config =
                        TrackEventConfig(
                            // currently `tracing:tracing-perfetto` records all events as part of
                            // the "rendering" category. In the future we should consider only
                            // setting this if sdk tracing is requested.
                            enabled_categories = listOf("rendering"),
                            disabled_categories = listOf("*"),
                        ),
                )
            ),
        )
    if (stackSamplingConfig != null) {
        dataSources += stackSamplingSource(config = stackSamplingConfig)
    }
    return configOf(dataSources)
}

internal fun TraceConfig.validateAndEncode(): ByteArray {
    val ftraceConfig = data_sources.firstNotNullOf { it.config?.ftrace_config }

    // check tags against known-supported tags based on SDK_INT / root status
    val supportedTags =
        AtraceTag.supported(api = Build.VERSION.SDK_INT, rooted = Shell.isSessionRooted())
            .map { it.tag }
            .toSet()

    val unsupportedTags = (ftraceConfig.atrace_categories - supportedTags)
    check(unsupportedTags.isEmpty()) {
        "Error - attempted to use unsupported atrace tags: $unsupportedTags"
    }

    if (Build.VERSION.SDK_INT < 28) {
        check(!ftraceConfig.atrace_apps.contains("*")) {
            "Support for wildcard (*) app matching in atrace added in API 28"
        }
    }

    if (Build.VERSION.SDK_INT < 24) {
        val packageList = ftraceConfig.atrace_apps.joinToString(",")
        check(packageList.length <= 91) {
            "Unable to trace package list (\"$packageList\").length = " +
                "${packageList.length} > 91 chars, which is the limit before API 24"
        }
    }
    return encode()
}
